#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Qwen模型微调参数解析模块
将参数解析逻辑从主训练脚本中分离出来
"""

import argparse
import os
import sys
import glob
import time
import re

# 获取项目根目录并添加到系统路径
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(PROJECT_ROOT)
from utils.model_utils import find_local_model_path


def find_latest_checkpoint(output_dir):
    """
    查找指定目录下最新的检查点（基于文件或目录的修改时间）
    
    Args:
        output_dir (str): 模型输出目录路径
        
    Returns:
        str or None: 最新检查点路径，如果没有找到则返回None
    """
    if not os.path.exists(output_dir):
        return None
        
    # 查找所有checkpoint-*目录
    checkpoint_dirs = glob.glob(os.path.join(output_dir, "checkpoint-*"))
    if not checkpoint_dirs:
        return None
    
    # 根据目录的修改时间确定最新的检查点
    def get_checkpoint_mtime(checkpoint_path):
        # 获取目录的最新修改时间
        try:
            return os.path.getmtime(checkpoint_path)
        except OSError:
            # 如果无法获取时间信息，则返回0
            return 0
    
    # 按修改时间排序并返回最新的检查点
    latest_checkpoint = max(checkpoint_dirs, key=get_checkpoint_mtime)
    return latest_checkpoint


def extract_checkpoint_number(checkpoint_path):
    """
    从检查点路径中提取编号（只从单个最新的检查点中提取）
    
    Args:
        checkpoint_path (str): 检查点路径
        
    Returns:
        int: 检查点编号，如果无法提取则返回0
    """
    if not checkpoint_path:
        return 0
    
    # 使用正则表达式提取checkpoint-后面的数字
    match = re.search(r'checkpoint-(\d+)', checkpoint_path)
    if match:
        return int(match.group(1))
    return 0


def get_finetune_args():
    """
    构建并返回微调参数解析器
    
    Returns:
        argparse.Namespace: 解析后的参数
    """
    parser = argparse.ArgumentParser(description="微调Qwen模型 - 增强版")
    parser.add_argument(
        "--model_path",
        default=find_local_model_path(PROJECT_ROOT),
        help="模型路径"
    )
    parser.add_argument(
        "--data_path",
        required=True,
        help="训练数据路径 (JSONL格式)"
    )
    parser.add_argument(
        "--output_dir",
        required=True,
        help="输出目录"
    )
    parser.add_argument(
        "--use_lora",
        action="store_true",
        help="是否使用LoRA微调"
    )
    parser.add_argument(
        "--lora_rank",
        type=int,
        default=8,
        help="LoRA秩参数"
    )
    parser.add_argument(
        "--max_length",
        type=int,
        default=512,
        help="最大序列长度"
    )
    parser.add_argument(
        "--num_train_epochs",
        type=int,
        default=10,
        help="训练轮数"
    )
    parser.add_argument(
        "--per_device_train_batch_size",
        type=int,
        default=4,
        help="每个设备的批次大小"
    )
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=2e-4,
        help="学习率"
    )
    parser.add_argument(
        "--use_cpu",
        action="store_true",
        help="是否使用CPU训练"
    )
    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=8,
        help="梯度累积步数"
    )
    # 新增参数
    parser.add_argument(
        "--resume_from_checkpoint",
        type=str,
        help="从检查点恢复训练（指定检查点路径）"
    )
    parser.add_argument(
        "--old_data_path", 
        type=str,
        help="旧数据路径（用于经验回放防遗忘）"
    )
    parser.add_argument(
        "--replay_ratio",
        type=float,
        default=0.0,
        help="经验回放比例 (0.0-1.0)"
    )
    parser.add_argument(
        "--version",
        type=str,
        help="指定版本号（不指定则自动生成）"
    )
    
    # 解析已知参数，以便在后续处理中使用output_dir
    args, _ = parser.parse_known_args()
    
    # 如果用户没有指定resume_from_checkpoint，则尝试查找output_dir下的最新检查点
    if not args.resume_from_checkpoint and args.output_dir:
        detected_checkpoint = find_latest_checkpoint(args.output_dir)
        if detected_checkpoint:
            print(f"🔍 自动检测到最新检查点: {detected_checkpoint}")
            # 将检测到的检查点路径添加到sys.argv中，确保在最终解析时能被正确识别
            sys.argv.append("--resume_from_checkpoint")
            sys.argv.append(detected_checkpoint)
            
            # 提取检查点编号并调整训练轮数
            checkpoint_number = extract_checkpoint_number(detected_checkpoint)
            if checkpoint_number > 0:
                # 将检查点编号添加到训练轮数参数中
                sys.argv.append("--num_train_epochs")
                # 计算数据条数除以10向上取整作为除数
                import json
                import math
                with open(args.data_path, 'r', encoding='utf-8') as f:
                    data_lines = sum(1 for line in f)
                divisor = max(1, math.ceil(data_lines / 10))  # 除数为数据条数除以10向上取整
                adjusted_epochs = math.ceil((args.num_train_epochs + checkpoint_number) / divisor)
                sys.argv.append(str(adjusted_epochs))
                print(f"🔄 调整训练轮数: ({args.num_train_epochs} + {checkpoint_number}) / {divisor}(数据条数:{data_lines}/10向上取整) = {adjusted_epochs}")
    
    # 重新解析所有参数
    return parser.parse_args()


if __name__ == "__main__":
    # 可以直接运行此脚本测试参数解析
    args = get_finetune_args()
    print("参数解析成功:")
    for arg, value in vars(args).items():
        print(f"  {arg}: {value}")